{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8ZekaWhZH2SX"
      },
      "outputs": [],
      "source": [
        "# @title ###### Licensed to the Apache Software Foundation (ASF), Version 2.0 (the \"License\")\n",
        "\n",
        "# Licensed to the Apache Software Foundation (ASF) under one\n",
        "# or more contributor license agreements. See the NOTICE file\n",
        "# distributed with this work for additional information\n",
        "# regarding copyright ownership. The ASF licenses this file\n",
        "# to you under the Apache License, Version 2.0 (the\n",
        "# \"License\"); you may not use this file except in compliance\n",
        "# with the License. You may obtain a copy of the License at\n",
        "#\n",
        "#   http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing,\n",
        "# software distributed under the License is distributed on an\n",
        "# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n",
        "# KIND, either express or implied. See the License for the\n",
        "# specific language governing permissions and limitations\n",
        "# under the License"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K6-p-DVrIFTY"
      },
      "source": [
        "# Vector Embedding Ingestion with Apache Beam and CloudSQL MySQL\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/apache/beam/blob/master/examples/notebooks/beam-ml/cloudsql_mysql_product_catalog_embeddings.ipynb\"><img src=\"https://raw.githubusercontent.com/google/or-tools/main/tools/colab_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/apache/beam/blob/master/examples/notebooks/beam-ml/cloudsql_mysql_product_catalog_embeddings.ipynb\"><img src=\"https://raw.githubusercontent.com/google/or-tools/main/tools/github_32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WWwFCLRHZPm4"
      },
      "source": [
        "# Introduction\n",
        "\n",
        "This Colab demonstrates how to generate embeddings from data and ingest them into [CloudSQL MySQL](https://cloud.google.com/sql/docs/mysql). We'll use Apache Beam and Dataflow for scalable data processing.\n",
        "\n",
        "The goal of this notebook is to make it easy for users to get started with generating embeddings at scale using Apache Beam and storing them in CloudSQL MySQL. We focus on building efficient ingestion pipelines that can handle various data sources and embedding models.\n",
        "\n",
        "## Example: Furniture Product Catalog\n",
        "\n",
        "We'll work with a sample e-commerce dataset representing a furniture product catalog. Each product has:\n",
        "\n",
        "*   **Structured fields:** `id`, `name`, `category`, `price`\n",
        "*   **Detailed text descriptions:** Longer text describing the product's features.\n",
        "*   **Additional metadata:** `material`, `dimensions`\n",
        "\n",
        "## Pipeline Overview\n",
        "We will build a pipeline to:\n",
        "1. Read product data\n",
        "2. Convert unstructured product data, to `Chunk`<sup>[1]</sup> type\n",
        "2. Generate Embeddings: Use a pre-trained Hugging Face model (via MLTransform) to create vector embeddings\n",
        "3. Write to CloudSQL MySQL: Store the embeddings in a CloudSQL MySQL vector database\n",
        "\n",
        "Here's a visualization of the data flow:\n",
        "\n",
        "| Stage                     | Data Representation                                      | Notes                                                                                                                   |\n",
        "| :------------------------ | :------------------------------------------------------- | :---------------------------------------------------------------------------------------------------------------------- |\n",
        "| **1. Ingest Data**      | `{`<br> `  \"id\": \"desk-001\",`<br> `  \"name\": \"Modern Desk\",`<br> `  \"description\": \"Sleek...\",`<br> `  \"category\": \"Desks\",`<br> `  ...`<br> `}` | Supports:<br>- Reading from batch (e.g., files, databases)<br>- Streaming sources (e.g., Pub/Sub).                 |\n",
        "| **2. Convert to Chunks** | `Chunk(` <br>  &nbsp;&nbsp;`id=\"desk-001\",` <br>  &nbsp;&nbsp;`content=Content(` <br>   &nbsp;&nbsp;&nbsp;&nbsp;`text=\"Modern Desk\"` <br> &nbsp;&nbsp; `),` <br>  &nbsp;&nbsp;`metadata={...}` <br> `)`       | - `Chunk` is the structured input for generating and ingesting embeddings.<br>- `chunk.content.text` is the field that is embedded.<br> - Converting to `Chunk` does not mean breaking data into smaller pieces,<br>&nbsp;&nbsp; it's simply organizing your data in a standard format for the embedding pipeline.<br> - `Chunk` allows data to flow seamlessly throughout embedding pipelines. |\n",
        "| **3. Generate Embeddings**| `Chunk(` <br>  &nbsp;&nbsp;`id=\"desk-001\",`<br>  &nbsp;&nbsp;`embedding=[-0.1, 0.6, ...],`<br>  `...)`  | Supports:<br>- Local Hugging Face models<br>- Remote Vertex AI models<br>- Custom embedding implementations.          |\n",
        "| **4. Write to CloudSQL MySQL** | **CloudSQL MySQL Table (Example Row):**<br>`id: desk-001`<br>`embedding: [-0.1, 0.6, ...]`<br> `name = \"Modern Desk\"`,<br>`Other fields ...` | Supports:<br>- Custom schemas<br>- Conflict resolution strategies for handling updates                               |\n",
        "\n",
        "\n",
        "[1]: Chunk represents an embeddable unit of input. It specifies which fields should be embedded and which fields should be treated as metadata. Converting to Chunk does not necessarily mean breaking your text into smaller pieces - it's primarily about structuring your data for the embedding pipeline. For very long texts that exceed the embedding model's maximum input size, you can optionally [use Langchain TextSplitters](https://beam.apache.org/releases/pydoc/2.63.0/apache_beam.ml.rag.chunking.langchain.html) to break the text into smaller `Chunk`'s.\n",
        "\n",
        "## Execution Environments\n",
        "\n",
        "This notebook demonstrates two execution environments:\n",
        "\n",
        "1. **DirectRunner (Local Execution)**: All examples in this notebook run on DirectRunner by default, which executes the pipeline locally. This is ideal for development, testing, and processing small datasets.\n",
        "\n",
        "2. **DataflowRunner (Distributed Execution)**: The [Run on Dataflow](#scrollTo=Quick_Start_Run_on_Dataflow) section demonstrates how to execute the same pipeline on Google Cloud Dataflow for scalable, distributed processing. This is recommended for production workloads and large datasets.\n",
        "\n",
        "All examples in this notebook can be adapted to run on Dataflow by following the pattern shown in the \"Run on Dataflow\" section."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z2eAyRECIP3z"
      },
      "source": [
        "# Connecting Apache Beam to CloudSQL MySQL\n",
        "\n",
        "Beam uses the [CloudSQL MySQL Java Connector](https://github.com/GoogleCloudPlatform/cloud-sql-jdbc-socket-factory/blob/main/docs/jdbc.md) to securely establish a connection to your database. Apache Beam supports any parameters that can be passed to the Java Connector e.g. IP types.\n",
        "\n",
        "# Setup and Prerequisites\n",
        "\n",
        "This example requires:\n",
        "1. A CloudSQL MySQL instance with [cloudsql_vector](https://cloud.google.com/sql/docs/mysql/vector-search#requirements) flag enabled\n",
        "2. Apache Beam 2.67.0 or later\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WhOOPUBa6PyW"
      },
      "source": [
        "## Install Packages and Dependencies\n",
        "\n",
        "First, let's install the Python packages required for the embedding and ingestion pipeline:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gCWRw2YE11wN"
      },
      "outputs": [],
      "source": [
        "# Apache Beam with GCP support\n",
        "!pip install apache_beam[interactive,gcp]>=2.67.0 --quiet\n",
        "# Huggingface sentence-transformers for embedding models\n",
        "!pip install sentence-transformers --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2FlMPmA0IUuv"
      },
      "outputs": [],
      "source": [
        "!pip show apache-beam"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4aqYZ_pG1oYb"
      },
      "source": [
        "Next, let's install cloud-sql-python-connector to help set up our test database."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eOYjnVDR87IE"
      },
      "outputs": [],
      "source": [
        "!pip install \"cloud-sql-python-connector[pymysql]>=1.0.0,<2.0.0\" sqlalchemy --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VhgbpTKzI-zI"
      },
      "source": [
        "## Database Setup\n",
        "\n",
        "To connect to CloudSQL MySQL, you'll need:\n",
        "1. GCP project ID where the CloudSQL MySQL instance is located\n",
        "2. The CloudSQL MySQL connection URI. This is the fully qualified connection name of the CloudSQL MySQL instance found in the google cloud console under CloudSQL > Instances > Instance > Connect to this Instance > Connection name.\n",
        "3. Database name. This is the name of the mysql database within your CloudSQL MySQL instance. The default database name is mysql.\n",
        "4. Database credentials\n",
        "5. A CloudSQL MySQL instance with cloudsql_vector flag enabled\n",
        "\n",
        "Replace these placeholder values with your actual CloudSQL MySQL connection details:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oqKQT0c_JB5f"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"\" # @param {type:'string'}\n",
        "\n",
        "CONNECTION_NAME = \"\" # @param {type:'string'}\n",
        "\n",
        "DB_NAME = \"\" #  @param {type:'string'}\n",
        "\n",
        "DB_USER = \"\" # @param {type:'string'}\n",
        "\n",
        "DB_PASSWORD = \"\" # @param {type:'string'}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "doK840yZZNdl"
      },
      "source": [
        "## Authenticate to Google Cloud\n",
        "\n",
        "To connect to the CloudSQL MySQL instance via the language conenctor, we authenticate with Google Cloud."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CLM12rbiZHTN"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "if 'google.colab' in sys.modules:\n",
        "  from google.colab import auth\n",
        "  auth.authenticate_user(project_id=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l_BBCKl7KKcb"
      },
      "outputs": [],
      "source": [
        " # @title SQLAlchemy + CloudSQL MySQL Connector helpers for creating tables and verifying data\n",
        "\n",
        "import sqlalchemy\n",
        "from sqlalchemy import text\n",
        "from sqlalchemy.exc import SQLAlchemyError\n",
        "from google.cloud.sql.connector import Connector\n",
        "\n",
        "def get_db_engine(connection_name: str, user: str, password: str, db: str, **connect_kwargs) -> sqlalchemy.engine.Engine:\n",
        "    \"\"\"\n",
        "    Creates a SQLAlchemy engine configured for CloudSQL MySQL.\n",
        "\n",
        "    To use this function, you may need to install necessary libraries:\n",
        "    'pip install google-cloud-sql-connector[pymysql] sqlalchemy'\n",
        "\n",
        "    Args:\n",
        "        connection_name: CloudSQL MySQL instance connection name (e.g., \"project:region:instance\").\n",
        "        user: The database user.\n",
        "        password: The database password.\n",
        "        db: The name of the database.\n",
        "        connect_kwargs: Additional keyword arguments for the connector (e.g., ip_type=\"PUBLIC\").\n",
        "\n",
        "    Returns:\n",
        "        A SQLAlchemy engine instance.\n",
        "    \"\"\"\n",
        "    connector = Connector()\n",
        "\n",
        "    def get_conn() -> sqlalchemy.engine.base.Connection:\n",
        "        \"\"\"Helper function to create a database connection.\"\"\"\n",
        "        conn = connector.connect(\n",
        "            connection_name,\n",
        "            \"pymysql\",  # Use the PyMySQL driver for MySQL\n",
        "            user=user,\n",
        "            password=password,\n",
        "            db=db,\n",
        "            **connect_kwargs\n",
        "        )\n",
        "        return conn\n",
        "\n",
        "    # Create the SQLAlchemy engine using the connection function\n",
        "    engine = sqlalchemy.create_engine(\n",
        "        \"mysql+pymysql://\",  # Use the MySQL+PyMySQL dialect\n",
        "        creator=get_conn,\n",
        "    )\n",
        "\n",
        "    # This hook ensures the connector is closed when the engine is disposed\n",
        "    engine.pool.dispose = lambda: connector.close()\n",
        "\n",
        "    return engine\n",
        "\n",
        "def setup_db_table_sqlalchemy(connection_name: str,\n",
        "                                 database: str,\n",
        "                                 table_name: str,\n",
        "                                 table_schema: str,\n",
        "                                 user: str,\n",
        "                                 password: str,\n",
        "                                 **connect_kwargs):\n",
        "    \"\"\"\n",
        "    Sets up a CloudSQL MySQL table using SQLAlchemy.\n",
        "\n",
        "    This function will drop the table if it already exists and then create it\n",
        "    based on the provided schema.\n",
        "\n",
        "    Args:\n",
        "        connection_name: CloudSQL MySQL instance connection name.\n",
        "        database: The name of the database.\n",
        "        table_name: The name of the table to create.\n",
        "        table_schema: SQL string defining the table columns. For MySQL, use types like\n",
        "                      'INT AUTO_INCREMENT PRIMARY KEY'. For embeddings, consider using\n",
        "                      'JSON' or 'BLOB' to store the vector data.\n",
        "                      Example: \"id INT AUTO_INCREMENT PRIMARY KEY, embedding JSON\"\n",
        "        user: The database user.\n",
        "        password: The database password.\n",
        "        connect_kwargs: Additional keyword arguments for the connector.\n",
        "    \"\"\"\n",
        "    engine = None\n",
        "    try:\n",
        "        engine = get_db_engine(connection_name, user, password, database, **connect_kwargs)\n",
        "\n",
        "        with engine.connect() as connection:\n",
        "            # Use autocommit for DDL statements\n",
        "            with connection.execution_options(isolation_level=\"AUTOCOMMIT\"):\n",
        "                print(\"Connected to MySQL DB successfully via SQLAlchemy!\")\n",
        "\n",
        "                # Use backticks for table names for MySQL compatibility\n",
        "                print(f\"Dropping table `{table_name}` if it exists...\")\n",
        "                connection.execute(text(f\"DROP TABLE IF EXISTS `{table_name}`;\"))\n",
        "\n",
        "                print(f\"Creating table `{table_name}`...\")\n",
        "                create_sql = f\"\"\"\n",
        "                CREATE TABLE IF NOT EXISTS `{table_name}` (\n",
        "                    {table_schema}\n",
        "                );\n",
        "                \"\"\"\n",
        "                connection.execute(text(create_sql))\n",
        "\n",
        "        print(\"MySQL table setup completed successfully!\")\n",
        "\n",
        "    except SQLAlchemyError as e:\n",
        "        print(f\"An SQLAlchemy error occurred during setup: {e}\")\n",
        "    except Exception as e:\n",
        "        print(f\"An unexpected error occurred during setup: {e}\")\n",
        "    finally:\n",
        "        if engine:\n",
        "            engine.dispose()\n",
        "\n",
        "def test_db_connection_sqlalchemy(connection_name: str,\n",
        "                                     database: str,\n",
        "                                     table_name: str,\n",
        "                                     user: str,\n",
        "                                     password: str,\n",
        "                                     **connect_kwargs):\n",
        "    \"\"\"\n",
        "    Tests the CloudSQL MySQL connection and verifies table existence.\n",
        "\n",
        "    Args:\n",
        "        connection_name: CloudSQL MySQL instance connection name.\n",
        "        database: The name of the database.\n",
        "        table_name: The name of the table to check for.\n",
        "        user: The database user.\n",
        "        password: The database password.\n",
        "        connect_kwargs: Additional keyword arguments for the connector.\n",
        "    \"\"\"\n",
        "    engine = None\n",
        "    try:\n",
        "        engine = get_db_engine(connection_name, user, password, database, **connect_kwargs)\n",
        "\n",
        "        with engine.connect() as connection:\n",
        "            print(\"Testing MySQL connection...\")\n",
        "            connection.execute(text(\"SELECT 1\"))\n",
        "            print(\"✓ Connection successful\")\n",
        "\n",
        "            # Check if table exists using information_schema.\n",
        "            # In MySQL, schema is the database, which can be found with DATABASE().\n",
        "            table_exists_query = text(\"\"\"\n",
        "                SELECT EXISTS (\n",
        "                    SELECT 1\n",
        "                    FROM information_schema.tables\n",
        "                    WHERE table_schema = DATABASE() AND table_name = :tname\n",
        "                );\n",
        "            \"\"\")\n",
        "            table_exists = connection.execute(table_exists_query, {\"tname\": table_name}).scalar()\n",
        "\n",
        "            if table_exists:\n",
        "                print(f\"✓ Table `{table_name}` exists in database `{database}`.\")\n",
        "            else:\n",
        "                print(f\"✗ Table `{table_name}` does NOT exist in database `{database}`.\")\n",
        "\n",
        "    except SQLAlchemyError as e:\n",
        "        print(f\"Connection test failed (SQLAlchemy error): {e}\")\n",
        "    except Exception as e:\n",
        "        print(f\"Connection test failed (Unexpected error): {e}\")\n",
        "    finally:\n",
        "        if engine:\n",
        "            engine.dispose()\n",
        "\n",
        "def verify_embeddings_sqlalchemy(connection_name: str,\n",
        "                                 database: str,\n",
        "                                 table_name: str,\n",
        "                                 user: str,\n",
        "                                 password: str,\n",
        "                                 embedding_column: str = \"embedding\",\n",
        "                                 **connect_kwargs):\n",
        "    \"\"\"\n",
        "    Connects to a CloudSQL MySQL table and prints all of its rows.\n",
        "\n",
        "    Args:\n",
        "        connection_name: CloudSQL MySQL instance connection name.\n",
        "        database: The name of the database.\n",
        "        table_name: The name of the table to query.\n",
        "        user: The database user.\n",
        "        password: The database password.\n",
        "        connect_kwargs: Additional keyword arguments for the connector.\n",
        "    \"\"\"\n",
        "    engine = None\n",
        "    try:\n",
        "        engine = get_db_engine(connection_name, user, password, database, **connect_kwargs)\n",
        "\n",
        "        with engine.connect() as connection:\n",
        "            # Use backticks for the table name for MySQL best practice\n",
        "            column_query = text(f\"\"\"\n",
        "                SELECT COLUMN_NAME\n",
        "                FROM INFORMATION_SCHEMA.COLUMNS\n",
        "                WHERE table_schema = :db_name\n",
        "                  AND table_name = :t_name\n",
        "                  AND COLUMN_NAME != '{embedding_column}'\n",
        "            \"\"\")\n",
        "\n",
        "            column_result = connection.execute(\n",
        "                column_query,\n",
        "                {\"db_name\": database, \"t_name\": table_name}\n",
        "            )\n",
        "\n",
        "            columns_to_select = [row[0] for row in column_result]\n",
        "\n",
        "            if not columns_to_select:\n",
        "                print(f\"No columns to display in `{table_name}` (after excluding '{embedding_column}').\")\n",
        "                return\n",
        "\n",
        "            # Construct the SELECT statement with the filtered columns, quoting them for safety\n",
        "            select_columns_str = \", \".join([f\"`{col}`\" for col in columns_to_select])\n",
        "            select_query = text(f\"SELECT {select_columns_str}, vector_to_string({embedding_column}) as {embedding_column} FROM `{table_name}`;\")\n",
        "\n",
        "            # Execute the query to get the data\n",
        "            result = connection.execute(select_query)\n",
        "            rows = result.mappings().all()\n",
        "\n",
        "            print(f\"\\nFound {len(rows)} rows in `{table_name}` (excluding '{embedding_column}' column):\")\n",
        "            print(\"-\" * 80)\n",
        "\n",
        "            if not rows:\n",
        "                print(\"Table is empty.\")\n",
        "            else:\n",
        "                # result.keys() will have the correct column names from the executed query\n",
        "                columns = result.keys()\n",
        "                for row in rows:\n",
        "                    for col in columns:\n",
        "                        print(f\"{col}: {row[col]}\")\n",
        "                    print(\"-\" * 80)\n",
        "    except SQLAlchemyError as e:\n",
        "        # Check specifically for ProgrammingError if the table might not exist\n",
        "        if isinstance(e, sqlalchemy.exc.ProgrammingError):\n",
        "            print(f\"Failed to query table `{table_name}`. Does it exist? Error: {e}\")\n",
        "        else:\n",
        "            print(f\"Failed to verify data (SQLAlchemy error): {e}\")\n",
        "    except Exception as e:\n",
        "        print(f\"Failed to verify data (Unexpected error): {e}\")\n",
        "    finally:\n",
        "        if engine:\n",
        "            engine.dispose()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "70z2O4nbOuaM"
      },
      "source": [
        "## Create Sample Product Catalog Data\n",
        "\n",
        "We'll create a typical e-commerce catalog where you might want to:\n",
        "- Generate embeddings for product text\n",
        "- Store vectors alongside product data\n",
        "- Enable vector similarity features\n",
        "\n",
        "Example product:\n",
        "```python\n",
        "{\n",
        "    \"id\": \"desk-001\",\n",
        "    \"name\": \"Modern Minimalist Desk\",\n",
        "    \"description\": \"Sleek minimalist desk with clean lines and a spacious work surface. \"\n",
        "                  \"Features cable management system and sturdy steel frame. \"\n",
        "                  \"Perfect for contemporary home offices and workspaces.\",\n",
        "    \"category\": \"Desks\",\n",
        "    \"price\": 399.99,\n",
        "    \"material\": \"Engineered Wood, Steel\",\n",
        "    \"dimensions\": \"60W x 30D x 29H inches\"\n",
        "}\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7_J__S8JOwJ_"
      },
      "outputs": [],
      "source": [
        "#@title Create sample data\n",
        "PRODUCTS_DATA = [\n",
        "    {\n",
        "        \"id\": \"desk-001\",\n",
        "        \"name\": \"Modern Minimalist Desk\",\n",
        "        \"description\": \"Sleek minimalist desk with clean lines and a spacious work surface. \"\n",
        "                      \"Features cable management system and sturdy steel frame. \"\n",
        "                      \"Perfect for contemporary home offices and workspaces.\",\n",
        "        \"category\": \"Desks\",\n",
        "        \"price\": 399.99,\n",
        "        \"material\": \"Engineered Wood, Steel\",\n",
        "        \"dimensions\": \"60W x 30D x 29H inches\"\n",
        "    },\n",
        "    {\n",
        "        \"id\": \"chair-001\",\n",
        "        \"name\": \"Ergonomic Mesh Office Chair\",\n",
        "        \"description\": \"Premium ergonomic office chair with breathable mesh back, \"\n",
        "                      \"adjustable lumbar support, and 4D armrests. Features synchronized \"\n",
        "                      \"tilt mechanism and memory foam seat cushion. Ideal for long work hours.\",\n",
        "        \"category\": \"Office Chairs\",\n",
        "        \"price\": 299.99,\n",
        "        \"material\": \"Mesh, Metal, Premium Foam\",\n",
        "        \"dimensions\": \"26W x 26D x 48H inches\"\n",
        "    },\n",
        "    {\n",
        "        \"id\": \"sofa-001\",\n",
        "        \"name\": \"Contemporary Sectional Sofa\",\n",
        "        \"description\": \"Modern L-shaped sectional with chaise lounge. Upholstered in premium \"\n",
        "                      \"performance fabric. Features deep seats, plush cushions, and solid \"\n",
        "                      \"wood legs. Perfect for modern living rooms.\",\n",
        "        \"category\": \"Sofas\",\n",
        "        \"price\": 1299.99,\n",
        "        \"material\": \"Performance Fabric, Solid Wood\",\n",
        "        \"dimensions\": \"112W x 65D x 34H inches\"\n",
        "    },\n",
        "    {\n",
        "        \"id\": \"table-001\",\n",
        "        \"name\": \"Rustic Dining Table\",\n",
        "        \"description\": \"Farmhouse-style dining table with solid wood construction. \"\n",
        "                      \"Features distressed finish and trestle base. Seats 6-8 people \"\n",
        "                      \"comfortably. Perfect for family gatherings.\",\n",
        "        \"category\": \"Dining Tables\",\n",
        "        \"price\": 899.99,\n",
        "        \"material\": \"Solid Pine Wood\",\n",
        "        \"dimensions\": \"72W x 42D x 30H inches\"\n",
        "    },\n",
        "    {\n",
        "        \"id\": \"bed-001\",\n",
        "        \"name\": \"Platform Storage Bed\",\n",
        "        \"description\": \"Modern queen platform bed with integrated storage drawers. \"\n",
        "                      \"Features upholstered headboard and durable wood slat support. \"\n",
        "                      \"No box spring needed. Perfect for maximizing bedroom space.\",\n",
        "        \"category\": \"Beds\",\n",
        "        \"price\": 799.99,\n",
        "        \"material\": \"Engineered Wood, Linen Fabric\",\n",
        "        \"dimensions\": \"65W x 86D x 48H inches\"\n",
        "    }\n",
        "]\n",
        "print(f\"\"\"✓ Created PRODUCTS_DATA with {len(PRODUCTS_DATA)} records\"\"\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KUHPsWzQFKpL"
      },
      "source": [
        "## Importing Pipeline Components\n",
        "\n",
        "We import the following for configuring our embedding ingestion pipeline:\n",
        "- `apache_beam.ml.rag.types.Chunk`, the structured input for generating and ingesting embeddings\n",
        "- `apache_beam.ml.rag.ingestion.cloudsql.CloudSQLMySQLVectorWriterConfig` for configuring write behavior like schema mapping and conflict resolution\n",
        "- `apache_beam.ml.rag.ingestion.cloudsql.LanguageConnectorConfig` to connect using the [CloudSQL MySQL language connector](https://github.com/GoogleCloudPlatform/cloud-sql-jdbc-socket-factory/blob/main/docs/jdbc.md)\n",
        "- `apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteTransform` to perform the write step using CloudSQL MySQL configs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fFMjPZaelTi2"
      },
      "outputs": [],
      "source": [
        "# CloudSQL imports\n",
        "from apache_beam.ml.rag.ingestion.cloudsql import CloudSQLMySQLVectorWriterConfig\n",
        "from apache_beam.ml.rag.ingestion.cloudsql import LanguageConnectorConfig\n",
        "\n",
        "\n",
        "from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteTransform\n",
        "from apache_beam.ml.rag.types import Chunk, Content\n",
        "from apache_beam.ml.rag.embeddings.huggingface import HuggingfaceTextEmbeddings\n",
        "\n",
        "# Apache Beam core\n",
        "import apache_beam as beam\n",
        "from apache_beam.options.pipeline_options import PipelineOptions\n",
        "from apache_beam.ml.transforms.base import MLTransform\n",
        "\n",
        "# JDBC and MySQL utilities\n",
        "from apache_beam.ml.rag.ingestion.jdbc_common import WriteConfig\n",
        "from apache_beam.ml.rag.ingestion.mysql_common import ColumnSpecsBuilder, ConflictResolution"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FjUzsUtXzFof"
      },
      "source": [
        "# What's next?\n",
        "\n",
        "This colab covers several use cases that you can explore based on your needs after completing the Setup and Prerequisites:\n",
        "\n",
        "🔰 **New to vector embeddings?**\n",
        "- [Start with Quick Start](#scrollTo=Quick_Start_Basic_Vector_Ingestion)\n",
        "- Uses simple out-of-box schema\n",
        "- Perfect for initial testing\n",
        "\n",
        "🚀 **Need to scale to large datasets?**\n",
        "- [Go to Run on Dataflow](#scrollTo=Quick_Start_Run_on_Dataflow)\n",
        "- Learn how to execute the same pipeline at scale\n",
        "- Fully managed\n",
        "- Process large datasets efficiently\n",
        "\n",
        "🎯 **Have a specific schema?**\n",
        "- [Go to Custom Schema](#scrollTo=Custom_Schema_with_Column_Mapping)\n",
        "- Learn to use different column names\n",
        "- Map metadata to individual columns\n",
        "\n",
        "🔄 **Need to update embeddings?**\n",
        "- [Check out Updating Embeddings](#scrollTo=Update_Embeddings_and_Metadata_with_Conflict_Resolution)\n",
        "- Handle conflicts\n",
        "- Selective field updates\n",
        "\n",
        "🔗 **Need to generate and Store Embeddings for Existing CloudSQL MySQL Data??**\n",
        "- [See Database Integration](#scrollTo=Adding_Embeddings_to_Existing_Database_Records)\n",
        "- Read data from your CloudSQL MySQL table.\n",
        "- Generate embeddings for the relevant fields.\n",
        "- Update your table (or a related table) with the generated embeddings.\n",
        "\n",
        "🤖 **Want to use Google's AI models?**\n",
        "- [Try Vertex AI Embeddings](#scrollTo=Generate_Embeddings_with_VertexAI_Text_Embeddings)\n",
        "- Use Google's powerful embedding models\n",
        "- Seamlessly integrate with other Google Cloud services\n",
        "\n",
        "🔄 Need real-time embedding  updates?\n",
        "\n",
        "- [Try Streaming Embeddings from PubSub](#scrollTo=Streaming_Embeddings_Updates_from_PubSub)\n",
        "- Process continuous data streams\n",
        "- Update embeddings in real-time as information changes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pLEi3Z4wKMOX"
      },
      "source": [
        "<a name=\"quickstart\"></a>\n",
        "# Quick Start: Basic Vector Ingestion\n",
        "\n",
        "This section shows the simplest way to generate embeddings and store them in CloudSQL MySQL."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LWqEgqjQOcbA"
      },
      "source": [
        "## Create table with default schema\n",
        "\n",
        "Before running the pipeline, we need a table to store our embeddings:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "93YnjdJkFWOi"
      },
      "outputs": [],
      "source": [
        "table_name = \"default_product_embeddings\"\n",
        "table_schema = f\"\"\"\n",
        "  id VARCHAR(255) PRIMARY KEY,\n",
        "  embedding VECTOR(384) USING VARBINARY,\n",
        "  content text,\n",
        "  metadata JSON\n",
        "\"\"\"\n",
        "setup_db_table_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name,table_schema, DB_USER, DB_PASSWORD)\n",
        "test_db_connection_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name, DB_USER, DB_PASSWORD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DikTnoGbOioG"
      },
      "source": [
        "## Configure Pipeline Components\n",
        "\n",
        "Now define the components that control the pipeline behavior:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M8rVyZ6o-Nep"
      },
      "source": [
        "### Convert ingested product data to embeddable Chunks\n",
        "- Our data is ingested as product dictionaries\n",
        "- Embedding generation and ingestion processes `Chunks`\n",
        "- We convert each product dictionary to a `Chunk` to configure what text to embed and what to treat as metadata"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rm_IX5U6mP_r"
      },
      "outputs": [],
      "source": [
        "from typing import Dict, Any\n",
        "\n",
        "# The create_chunk function converts our product dictionaries to Chunks.\n",
        "# This doesn't split the text - it simply structures it in the format\n",
        "# expected by the embedding pipeline components.\n",
        "def create_chunk(product: Dict[str, Any]) -> Chunk:\n",
        "    \"\"\"Convert a product dictionary into a Chunk object.\n",
        "\n",
        "       The pipeline components (MLTransform, VectorDatabaseWriteTransform)\n",
        "       work with Chunk objects. This function:\n",
        "       1. Extracts text we want to embed\n",
        "       2. Preserves product data as metadata\n",
        "       3. Creates a Chunk in the expected format\n",
        "\n",
        "    Args:\n",
        "        product: Dictionary containing product information\n",
        "\n",
        "    Returns:\n",
        "        Chunk: A Chunk object ready for embedding\n",
        "    \"\"\"\n",
        "    return Chunk(\n",
        "        content=Content(\n",
        "            text=f\"{product['name']}: {product['description']}\"\n",
        "        ), # The text that will be embedded\n",
        "        id=product['id'],  # Use product ID as chunk ID\n",
        "        metadata=product,  # Store all product info in metadata\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xJaI9m3D7Vw-"
      },
      "source": [
        "### Generate embeddings with HuggingFace"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0dlm1fjQh2dX"
      },
      "source": [
        "We use a local pre-trained Hugging Face model to create vector embeddings from the product descriptions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E5LkHmjV7l2S"
      },
      "outputs": [],
      "source": [
        "huggingface_embedder = HuggingfaceTextEmbeddings(\n",
        "    model_name=\"sentence-transformers/all-MiniLM-L6-v2\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vVv8hD5wQo3w"
      },
      "source": [
        "### Write to CloudSQL MySQL\n",
        "\n",
        "The default CloudSQLMySQLVectorWriterConfig maps Chunk fields to database columns as:\n",
        "\n",
        "| Database Column | Chunk Field | Description |\n",
        "|----------------|-------------|-------------|\n",
        "| id             | chunk.id    | Unique identifier |\n",
        "| embedding      | chunk.embedding.dense_embedding | Vector representation |\n",
        "| content        | chunk.content.text | Text that was embedded |\n",
        "| metadata       | chunk.metadata | Additional data as JSONB |"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "moKsz_6xQt-E"
      },
      "outputs": [],
      "source": [
        "# Configure the language connector so we can connect securely\n",
        "connector_config = LanguageConnectorConfig(\n",
        "    username=DB_USER,\n",
        "    password=DB_PASSWORD,\n",
        "    database_name=DB_NAME,\n",
        "    instance_name=CONNECTION_NAME\n",
        ")\n",
        "cloudsql_writer_config = CloudSQLMySQLVectorWriterConfig(\n",
        "    connection_config=connector_config,\n",
        "    table_name=table_name\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ww2BPxTNKmL2"
      },
      "source": [
        "## Assemble and Run Pipeline\n",
        "\n",
        "Now we can create our pipeline that:\n",
        "1. Takes our product data\n",
        "2. Converts each product to a Chunk\n",
        "3. Generates embeddings for each Chunk\n",
        "4. Stores everything in CloudSQL MySQL"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lyS3IpNBDgYw"
      },
      "outputs": [],
      "source": [
        "import tempfile\n",
        "\n",
        "# Executing on DirectRunner (local execution)\n",
        "with beam.Pipeline() as p:\n",
        "    _ = (\n",
        "            p\n",
        "            | 'Create Products' >> beam.Create(PRODUCTS_DATA)\n",
        "            | 'Convert to Chunks' >> beam.Map(create_chunk)\n",
        "            | 'Generate Embeddings' >> MLTransform(write_artifact_location=tempfile.mkdtemp())\n",
        "              .with_transform(huggingface_embedder)\n",
        "            | 'Write to CloudSQL' >> VectorDatabaseWriteTransform(\n",
        "                cloudsql_writer_config\n",
        "            )\n",
        "        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qm97EAww6RvW"
      },
      "source": [
        "## Verify Embeddings\n",
        "Let's check what was written to our CloudSQL MySQL table:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-H3t2cIN6lO_"
      },
      "outputs": [],
      "source": [
        "verify_embeddings_sqlalchemy(connection_name=CONNECTION_NAME, database=DB_NAME, table_name=table_name, user=DB_USER, password=DB_PASSWORD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lz5itufZ31KB"
      },
      "source": [
        "## Quick Start Summary\n",
        "\n",
        "In this section, you learned how to:\n",
        "- Convert product data to the Chunk format expected by embedding pipelines\n",
        "- Generate embeddings using a HuggingFace model\n",
        "- Configure and run a basic embedding ingestion pipeline\n",
        "- Store embeddings and metadata in CloudSQL MySQL\n",
        "\n",
        "This basic pattern forms the foundation for all the advanced use cases covered in the following sections."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OqojLgpJKUGk"
      },
      "source": [
        "# Quick Start: Run on Dataflow\n",
        "\n",
        "This section demonstrates how to launch the Quick Start embedding pipeline on Google Cloud Dataflow from the colab. While previous examples used DirectRunner for local execution, Dataflow provides a fully managed, distributed execution environment that is:\n",
        "- Scalable: Automatically scales to handle large datasets\n",
        "- Fault-tolerant: Handles worker failures and ensures exactly-once processing\n",
        "- Fully managed: No need to provision or manage infrastructure\n",
        "\n",
        "For more in-depth documentation to package your pipeline into a python file and launch a DataFlow job from the command line see [Create Dataflow pipeline using Python](https://cloud.google.com/dataflow/docs/quickstarts/create-pipeline-python)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zrMJSm-JUVGY"
      },
      "source": [
        "## Create the CloudSQL MySQL table with default schema\n",
        "\n",
        "Before running the pipeline, we need a table to store our embeddings:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tgAvMT-yUixY"
      },
      "outputs": [],
      "source": [
        "table_name = \"default_dataflow_product_embeddings\"\n",
        "table_schema = f\"\"\"\n",
        "  id VARCHAR(255) PRIMARY KEY,\n",
        "  embedding VECTOR(384) USING VARBINARY,\n",
        "  content text,\n",
        "  metadata JSON\n",
        "\"\"\"\n",
        "setup_db_table_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name,table_schema, DB_USER, DB_PASSWORD)\n",
        "test_db_connection_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name, DB_USER, DB_PASSWORD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mcZATJbaOec0"
      },
      "source": [
        "## Save our Pipeline to a python file\n",
        "\n",
        "To launch our pipeline job on DataFlow, we\n",
        "1. Add command line arguments for passing pipeline options like CloudSQL MySQL credentioals\n",
        "2. Save our pipeline code to a local file `basic_ingestion_pipeline.py`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CzhIiBdqOknd"
      },
      "outputs": [],
      "source": [
        "file_content = \"\"\"\n",
        "import apache_beam as beam\n",
        "from apache_beam.options.pipeline_options import PipelineOptions\n",
        "import argparse\n",
        "import tempfile\n",
        "\n",
        "from apache_beam.ml.transforms.base import MLTransform\n",
        "from apache_beam.ml.rag.types import Chunk, Content\n",
        "from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteTransform\n",
        "from apache_beam.ml.rag.ingestion.cloudsql import CloudSQLMySQLVectorWriterConfig, LanguageConnectorConfig\n",
        "from apache_beam.ml.rag.embeddings.huggingface import HuggingfaceTextEmbeddings\n",
        "from apache_beam.options.pipeline_options import SetupOptions\n",
        "\n",
        "PRODUCTS_DATA = [\n",
        "    {\n",
        "        \"id\": \"desk-001\",\n",
        "        \"name\": \"Modern Minimalist Desk\",\n",
        "        \"description\": \"Sleek minimalist desk with clean lines and a spacious work surface. \"\n",
        "                      \"Features cable management system and sturdy steel frame. \"\n",
        "                      \"Perfect for contemporary home offices and workspaces.\",\n",
        "        \"category\": \"Desks\",\n",
        "        \"price\": 399.99,\n",
        "        \"material\": \"Engineered Wood, Steel\",\n",
        "        \"dimensions\": \"60W x 30D x 29H inches\"\n",
        "    },\n",
        "    {\n",
        "        \"id\": \"chair-001\",\n",
        "        \"name\": \"Ergonomic Mesh Office Chair\",\n",
        "        \"description\": \"Premium ergonomic office chair with breathable mesh back, \"\n",
        "                      \"adjustable lumbar support, and 4D armrests. Features synchronized \"\n",
        "                      \"tilt mechanism and memory foam seat cushion. Ideal for long work hours.\",\n",
        "        \"category\": \"Office Chairs\",\n",
        "        \"price\": 299.99,\n",
        "        \"material\": \"Mesh, Metal, Premium Foam\",\n",
        "        \"dimensions\": \"26W x 26D x 48H inches\"\n",
        "    }\n",
        "]\n",
        "\n",
        "def run(argv=None):\n",
        "    parser = argparse.ArgumentParser()\n",
        "    parser.add_argument(\n",
        "        '--connection_name',\n",
        "        required=True,\n",
        "        help='CloudSQL MySQL instance uri'\n",
        "    )\n",
        "    parser.add_argument(\n",
        "        '--cloudsql_database',\n",
        "        default='mysql',\n",
        "        help='CloudSQL MySQL database name'\n",
        "    )\n",
        "    parser.add_argument(\n",
        "        '--cloudsql_table',\n",
        "        required=True,\n",
        "        help='CloudSQL MySQL table name'\n",
        "    )\n",
        "    parser.add_argument(\n",
        "        '--cloudsql_username',\n",
        "        required=True,\n",
        "        help='CloudSQL MySQL user name'\n",
        "    )\n",
        "    parser.add_argument(\n",
        "        '--cloudsql_password',\n",
        "        required=True,\n",
        "        help='CloudSQL MySQL password'\n",
        "    )\n",
        "    known_args, pipeline_args = parser.parse_known_args(argv)\n",
        "\n",
        "    pipeline_options = PipelineOptions(pipeline_args)\n",
        "    pipeline_options.view_as(SetupOptions).save_main_session = True\n",
        "\n",
        "    with beam.Pipeline(options=pipeline_options) as p:\n",
        "        _ = (\n",
        "              p\n",
        "              | 'Create Products' >> beam.Create(PRODUCTS_DATA)\n",
        "              | 'Convert to Chunks' >> beam.Map(lambda product: Chunk(\n",
        "                    content=Content(\n",
        "                        text=f\"{product['name']}: {product['description']}\"\n",
        "                    ), # The text that will be embedded\n",
        "                    id=product['id'],  # Use product ID as chunk ID\n",
        "                    metadata=product,  # Store all product info in metadata\n",
        "                )\n",
        "              )\n",
        "              | 'Generate Embeddings' >> MLTransform(write_artifact_location=tempfile.mkdtemp())\n",
        "                .with_transform(\n",
        "                    HuggingfaceTextEmbeddings(model_name=\"sentence-transformers/all-MiniLM-L6-v2\")\n",
        "                )\n",
        "              | 'Write to CloudSQL MySQL' >> VectorDatabaseWriteTransform(\n",
        "                  CloudSQLMySQLVectorWriterConfig(\n",
        "                      connection_config=LanguageConnectorConfig(\n",
        "                        username=known_args.cloudsql_username,\n",
        "                        password=known_args.cloudsql_password,\n",
        "                        database_name=known_args.cloudsql_database,\n",
        "                        instance_name=known_args.connection_name\n",
        "                      ),\n",
        "                      table_name=known_args.cloudsql_table\n",
        "                  )\n",
        "              )\n",
        "          )\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    run()\n",
        "\"\"\"\n",
        "\n",
        "with open(\"basic_ingestion_pipeline.py\", \"w\") as f:\n",
        "    f.write(file_content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y_1IMXx7UuG4"
      },
      "source": [
        "## Authenticate with Google Cloud\n",
        "\n",
        "To launch a pipeline on Google Cloud, authenticate this notebook. Replace `<PROJECT_ID>` with your Google Cloud project ID"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WxrW-zlgRDLk"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"<project_id_>\" # @param {type:'string'}\n",
        "import os\n",
        "os.environ['PROJECT_ID'] = PROJECT_ID"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GswFBa10Qxkx"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "if 'google.colab' in sys.modules:\n",
        "  from google.colab import auth\n",
        "  auth.authenticate_user(project_id=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7sELV2KeRG2c"
      },
      "source": [
        "## Configure the Pipeline options\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nVDW0Q9iS_Pk"
      },
      "source": [
        "To run the pipeline on DataFlow we need\n",
        "- A gcs bucket for staging DataFlow files. Replace `<BUCKET_NAME>`: the name of a valid Google Cloud Storage bucket.\n",
        "- Optionally set the Google Cloud region that you want to run Dataflow in. Replace `<REGION>` with the desired location.\n",
        "- Optionally provide `NETWORK` and `SUBNETWORK` for dataflow workers to run on."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qxFJflLiTMua"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "BUCKET_NAME = '' # @param {type:'string'}\n",
        "REGION = 'us-central1' # @param {type:'string'}\n",
        "\n",
        "NETWORK = '' # @param {type:'string'}\n",
        "SUBNETWORK = '' # @param {type:'string'}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WWjjqwV-aJFi"
      },
      "source": [
        "## Provide additional Python dependencies to be installed on Worker VM's\n",
        "\n",
        "We are making use of the HuggingFace `sentence-transformers` package to generate embeddings. Since this package is not installed on Worker VM's by default, we create a requirements.txt file with the additional dependencies to be installed on worker VM's.\n",
        "\n",
        "See [Managing Python Pipeline Dependencies](https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/) for more details.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hkxmk6aTJSKW"
      },
      "outputs": [],
      "source": [
        "!echo \"sentence-transformers\" > ./requirements.txt\n",
        "!cat ./requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NXgGZeOsY2O7"
      },
      "source": [
        "## Run Pipeline on Dataflow\n",
        "\n",
        "We launch the pipeline via the command line, passing\n",
        "- CloudSQL MySQL pipeline arguments defined in `basic_ingestion_pipeline.py`\n",
        "- GCP Project ID\n",
        "- Job Region\n",
        "- The runner (DataflowRunner)\n",
        "- Temp and Staging GCS locations for Pipeline artifacts\n",
        "- Requirement file location for additional dependencies\n",
        "- (Optional) The VPC network and Subnetwork that has access to the CloudSQL MySQL instance\n",
        "\n",
        "Once the job is launched, you can monitor its progress in the Google Cloud Console:\n",
        "1. Go to https://console.cloud.google.com/dataflow/jobs\n",
        "2. Select your project\n",
        "3. Click on the job named \"cloudsql-dataflow-basic-embedding-ingest\"\n",
        "4. View detailed execution graphs, logs, and metrics"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fUeG_hEb5Qbb"
      },
      "outputs": [],
      "source": [
        "command_parts = [\n",
        "    \"python ./basic_ingestion_pipeline.py\",\n",
        "    f\"--project={PROJECT_ID}\",\n",
        "    f\"--cloudsql_username={DB_USER}\",\n",
        "    f\"--connection_name={CONNECTION_NAME}\",\n",
        "    f\"--cloudsql_password={DB_PASSWORD}\",\n",
        "    f\"--cloudsql_table=default_dataflow_product_embeddings\",\n",
        "    f\"--cloudsql_database={DB_NAME}\",\n",
        "    f\"--job_name=cloudsql-dataflow-basic-embedding-ingest\",\n",
        "    f\"--region={REGION}\",\n",
        "    \"--runner=DataflowRunner\",\n",
        "    f\"--temp_location=gs://{BUCKET_NAME}/temp\",\n",
        "    f\"--staging_location=gs://{BUCKET_NAME}/staging\",\n",
        "    \"--requirements_file=requirements.txt\",\n",
        "]\n",
        "\n",
        "if NETWORK:\n",
        "    command_parts.append(f\"--network={NETWORK}\")\n",
        "\n",
        "if SUBNETWORK:\n",
        "    command_parts.append(f\"--subnetwork=regions/{REGION}/subnetworks/{SUBNETWORK}\")\n",
        "\n",
        "final_command = \" \".join(command_parts)\n",
        "\n",
        "print(\"Generated command:\\n\", final_command)\n",
        "!{final_command}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sp_M6tJbWXTw"
      },
      "source": [
        "## Verify the Written Embeddings\n",
        "\n",
        "Let's check what was written to our CloudSQL MySQL table:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A11PeldtWXvP"
      },
      "outputs": [],
      "source": [
        "verify_embeddings_sqlalchemy(connection_name=CONNECTION_NAME, database=DB_NAME, table_name='default_dataflow_product_embeddings', user=DB_USER, password=DB_PASSWORD)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-2hTEi-jzYN6"
      },
      "source": [
        "# Advanced Use Cases\n",
        "\n",
        "This section demonstrates more complex scenarios for using CloudSQL MySQL with Apache Beam for vector embeddings.\n",
        "\n",
        "🎯 **Have a specific schema?**\n",
        "- [Go to Custom Schema](#scrollTo=Custom_Schema_with_Column_Mapping)\n",
        "- Learn to use different column names and transform values\n",
        "- Map metadata to individual columns\n",
        "\n",
        "🔄 **Need to update embeddings?**\n",
        "- [Check out Updating Embeddings](#scrollTo=Update_Embeddings_and_Metadata_with_Conflict_Resolution)\n",
        "- Handle conflicts\n",
        "- Selective field updates\n",
        "\n",
        "🔗 **Need to generate and Store Embeddings for Existing CloudSQL MySQL Data??**\n",
        "- [See Database Integration](#scrollTo=Adding_Embeddings_to_Existing_Database_Records)\n",
        "- Read data from your CloudSQL MySQL table.\n",
        "- Generate embeddings for the relevant fields.\n",
        "- Update your table (or a related table) with the generated embeddings.\n",
        "\n",
        "🤖 **Want to use Google's AI models?**\n",
        "- [Try Vertex AI Embeddings](#scrollTo=Generate_Embeddings_with_VertexAI_Text_Embeddings)\n",
        "- Use Google's powerful embedding models\n",
        "- Seamlessly integrate with other Google Cloud services\n",
        "\n",
        "🔄 Need real-time embedding  updates?\n",
        "\n",
        "- [Try Streaming Embeddings from PubSub](#scrollTo=Streaming_Embeddings_Updates_from_PubSub)\n",
        "- Process continuous data streams\n",
        "- Update embeddings in real-time as information changes\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qGaH_TqEzn8r"
      },
      "source": [
        "## Custom Schema with Column Mapping <a name=\"custom-schema\"></a>\n",
        "\n",
        "In this example, we'll create a custom schema that:\n",
        "- Uses different column names\n",
        "- Maps metadata to individual columns\n",
        "- Uses functions to transform values"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R4d9W6ry_CN8"
      },
      "source": [
        "### ColumnSpec and ColumnSpecsBuilder\n",
        "\n",
        "\n",
        "ColumnSpec specifies how to map data to a database column. For example:\n",
        "```python\n",
        "from apache_beam.ml.rag.ingestion.mysql_common import ColumnSpecsBuilder\n",
        "\n",
        "ColumnSpec(\n",
        "    column_name=\"price\",          # Database column\n",
        "    python_type=float,            # Python Type for the value\n",
        "    value_fn=lambda c: c.metadata['price'],  # Extract price from Chunk metadata to get actual value\n",
        "    placeholder=\"ROUND(?, 2)\"      # Optional SQL cast or function\n",
        ")\n",
        "```\n",
        "creates an INSERT statement like:\n",
        "```sql\n",
        "INSERT INTO table (price) VALUES (?::decimal)\n",
        "```\n",
        "where the `?` placeholder is poulated with the value from our ingested data.\n",
        "\n",
        "`ColumnSpecsBuilder` provides a builder and convenience methods to create these `ColumnSpecs`:\n",
        "\n",
        "1. Core Field Mapping\n",
        "   - `with_id_spec()` => Insert chunk.id as text in \"id\" column\n",
        "   - `with_embedding_spec()` => Insert chunk.embedding casted to `VECTOR` via `string_to_vector(?)` in \"embedding\" column\n",
        "   - `with_content_spec()` => Insert `chunk.content`.text as text in \"content\" column\n",
        "\n",
        "   Note: All `with_id_spec`, `with_embedding_spec`, etc. methods allow overriding `column_name`, `python_type`, and `value_fn`.\n",
        "\n",
        "2. Metadata Extraction\n",
        "   - `add_metadata_field`: Creates a column from a `chunk.metadata` field\n",
        "   - Handles type conversion based on specified SQL type\n",
        "\n",
        "3. Custom Fields\n",
        "   - `add_custom_column_spec`: Grants complete control over mapping `Chunk` data to database rows using `ColumnSpec`\n",
        "\n",
        "Now, lets the table to store our embeddings:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6XpYLcCu80Dy"
      },
      "source": [
        "### Create Custom Schema Table"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6bUd6vprzh7O"
      },
      "outputs": [],
      "source": [
        "table_name = \"custom_product_embeddings\"\n",
        "table_schema = \"\"\"\n",
        "    product_id VARCHAR(255) PRIMARY KEY,\n",
        "    vector_embedding VECTOR(384) USING VARBINARY,\n",
        "    product_name VARCHAR(255),\n",
        "    description TEXT,\n",
        "    price DECIMAL,\n",
        "    category VARCHAR(255),\n",
        "    display_text VARCHAR(255),\n",
        "    model_name VARCHAR(255),\n",
        "    created_at TIMESTAMP\n",
        "\"\"\"\n",
        "setup_db_table_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name,table_schema, DB_USER, DB_PASSWORD)\n",
        "test_db_connection_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name, DB_USER, DB_PASSWORD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ScCVCZFo-Fcv"
      },
      "source": [
        "### Configure Pipeline Components"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g9-f7tcf-0qC"
      },
      "source": [
        "#### Write to custom schema using ColumnSpecsBuilder\n",
        "\n",
        "\n",
        "We configure ConlumnSpecsBuilder to map data as:\n",
        "\n",
        "| Database Column | Chunk Field                               |\n",
        "|-----------------|-------------------------------------------|\n",
        "| `product_id`    | `chunk.id`                                |\n",
        "| `vector_embedding`| `chunk.embedding.dense_embedding`        |\n",
        "| `description`   | `chunk.content.text`                      |\n",
        "| `product_name`  | `chunk.metadata['name']`                  |\n",
        "| `price`         | `chunk.metadata['price']`                 |\n",
        "| `category`      | `chunk.metadata['category']`              |\n",
        "| `display_text`  | *Function that combines product name and price* |\n",
        "| `model_name`    | *Function that returns the model name: \"all-MiniLM-L6-v2\"* |\n",
        "| `created_at`    | *Function that returns the current timestamp cast to a SQL timestamp* |"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TAq6ydMn-5Uu"
      },
      "outputs": [],
      "source": [
        "from apache_beam.ml.rag.ingestion.mysql_common import ColumnSpecsBuilder\n",
        "from apache_beam.ml.rag.ingestion.mysql_common import ColumnSpec\n",
        "from datetime import datetime\n",
        "\n",
        "column_specs = (\n",
        "    ColumnSpecsBuilder()\n",
        "    # Write chunk.id to a column named \"product_id\"\n",
        "    .with_id_spec(column_name='product_id')\n",
        "    # Write chunk.embedding.dense_embedding to a column named \"vector_embedding\"\n",
        "    .with_embedding_spec(column_name='vector_embedding')\n",
        "    # Write chunk.content.text to a column named \"description\"\n",
        "    .with_content_spec(column_name='description')\n",
        "    # Write chunk.metadata.['product_name'] to a column named \"product_name\"\n",
        "    .add_metadata_field(\n",
        "        field='name',\n",
        "        column_name='product_name',\n",
        "        python_type=str\n",
        "    )\n",
        "    # Write chunk.metadata.['price'] to a column named \"price\"\n",
        "    .add_metadata_field(\n",
        "        field='price',\n",
        "        column_name='price',\n",
        "        python_type=float\n",
        "    )\n",
        "    # Write chunk.metadata.['category'] to a column named \"category\"\n",
        "    .add_metadata_field(\n",
        "        field='category',\n",
        "        column_name='category',\n",
        "        python_type=str\n",
        "    )\n",
        "    # Write custom field using value_fn to column named \"display_text\" using\n",
        "    # ColumnSpec.text convenience method\n",
        "    .add_custom_column_spec(\n",
        "        ColumnSpec.text(\n",
        "          column_name='display_text',\n",
        "        value_fn=lambda chunk: \\\n",
        "          f\"{chunk.metadata['name']} - ${chunk.metadata['price']:.2f}\"\n",
        "        )\n",
        "    )\n",
        "    # Store model used to generate embedding using ColumnSpec constructor\n",
        "    .add_custom_column_spec(\n",
        "        ColumnSpec(\n",
        "          column_name='model_name',\n",
        "          python_type=str,\n",
        "          value_fn=lambda _: \"all-MiniLM-L6-v2\"\n",
        "        )\n",
        "    )\n",
        "    .add_custom_column_spec(\n",
        "        ColumnSpec(\n",
        "          column_name='created_at',\n",
        "          python_type=str,\n",
        "          value_fn=lambda _: datetime.now().isoformat()\n",
        "        )\n",
        "    )\n",
        "    .build()\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MBfLVL6XX2mF"
      },
      "source": [
        "### Assemble and Run Pipeline\n",
        "\n",
        "Now we can create our pipeline that will:\n",
        "1. Take our product data\n",
        "2. Convert each product to a Chunk\n",
        "3. Generate embeddings for each Chunk\n",
        "4. Store everything in CloudSQL MySQL with our custom schema configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4V-ILUlWVVX8"
      },
      "outputs": [],
      "source": [
        "import tempfile # For storing MLTransform artifacts\n",
        "\n",
        "# Executing on DirectRunner (local execution)\n",
        "with beam.Pipeline() as p:\n",
        "    _ = (\n",
        "            p\n",
        "            | 'Create Products' >> beam.Create(PRODUCTS_DATA)\n",
        "            | 'Convert to Chunks' >> beam.Map(lambda product: Chunk(\n",
        "                    content=Content(\n",
        "                        text=f\"{product['name']}: {product['description']}\"\n",
        "                    ), # The text that will be embedded\n",
        "                    id=product['id'],  # Use product ID as chunk ID\n",
        "                    metadata=product,  # Store all product info in metadata\n",
        "                )\n",
        "              )\n",
        "            | 'Generate Embeddings' >> MLTransform(write_artifact_location=tempfile.mkdtemp())\n",
        "              .with_transform(HuggingfaceTextEmbeddings(model_name=\"sentence-transformers/all-MiniLM-L6-v2\"))\n",
        "            | 'Write to CloudSQL MySQL' >> VectorDatabaseWriteTransform(\n",
        "                CloudSQLMySQLVectorWriterConfig(\n",
        "                    connection_config=LanguageConnectorConfig(\n",
        "                        username=DB_USER,\n",
        "                        password=DB_PASSWORD,\n",
        "                        database_name=DB_NAME,\n",
        "                        instance_name=CONNECTION_NAME\n",
        "                    ),\n",
        "                    table_name=table_name,\n",
        "                    column_specs=column_specs\n",
        "                )\n",
        "            )\n",
        "        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LCpoJkBpYEsH"
      },
      "source": [
        "### Verify the Written Embeddings\n",
        "\n",
        "Let's check what was written to our CloudSQL MySQL table:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B2UDOZL0VZ-p"
      },
      "outputs": [],
      "source": [
        "verify_embeddings_sqlalchemy(connection_name=CONNECTION_NAME, database=DB_NAME, table_name=table_name, user=DB_USER, password=DB_PASSWORD, embedding_column=\"vector_embedding\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DQyJoyZic9GT"
      },
      "source": [
        "## Update Embeddings and Metadata with Conflict Resolution <a name=\"updating\"></a>\n",
        "\n",
        "This section demonstrates how to handle periodic updates to product descriptions and their embeddings using the default schema. We'll show how embeddings and metadata get updated when product descriptions change.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jwLHGKfNdbEG"
      },
      "source": [
        "### Create table with desired schema\n",
        "\n",
        "Let's use the same default schema as in Quick Start:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0vK-b4xkXtgJ"
      },
      "outputs": [],
      "source": [
        "table_name = \"mutable_product_embeddings\"\n",
        "table_schema = f\"\"\"\n",
        "  id VARCHAR(255) PRIMARY KEY,\n",
        "  embedding VECTOR(384) USING VARBINARY,\n",
        "  content text,\n",
        "  metadata JSON,\n",
        "  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\n",
        "\"\"\"\n",
        "setup_db_table_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name,table_schema, DB_USER, DB_PASSWORD)\n",
        "test_db_connection_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name, DB_USER, DB_PASSWORD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hhl2URWceSg_"
      },
      "source": [
        "### Sample Data: Day 1 vs Day 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t4z8tM_leZV8"
      },
      "outputs": [],
      "source": [
        "PRODUCTS_DATA_DAY1 = [\n",
        "    {\n",
        "        \"id\": \"desk-001\",\n",
        "        \"name\": \"Modern Minimalist Desk\",\n",
        "        \"description\": \"Sleek minimalist desk with clean lines and a spacious work surface. \"\n",
        "                      \"Features cable management system and sturdy steel frame.\",\n",
        "        \"category\": \"Desks\",\n",
        "        \"price\": 399.99,\n",
        "        \"update_timestamp\": \"2024-02-18\"\n",
        "    }\n",
        "]\n",
        "\n",
        "PRODUCTS_DATA_DAY2 = [\n",
        "    {\n",
        "        \"id\": \"desk-001\",  # Same ID as Day 1\n",
        "        \"name\": \"Modern Minimalist Desk\",\n",
        "        \"description\": \"Updated: Sleek minimalist desk with built-in wireless charging. \"\n",
        "                      \"Features cable management system, sturdy steel frame, and Qi charging pad. \"\n",
        "                      \"Perfect for modern tech-enabled workspaces.\",\n",
        "        \"category\": \"Smart Desks\",  # Category changed\n",
        "        \"price\": 449.99,  # Price increased\n",
        "        \"update_timestamp\": \"2024-02-19\"\n",
        "    }\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W_UTcRz9eskE"
      },
      "source": [
        "### Configure Pipeline Components"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PWvtwVmUedzw"
      },
      "source": [
        "#### Writer with Conflict Resolution"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y2XEwxw6ee4b"
      },
      "outputs": [],
      "source": [
        "from apache_beam.ml.rag.ingestion.cloudsql import (\n",
        "    CloudSQLMySQLVectorWriterConfig,\n",
        "    LanguageConnectorConfig,\n",
        ")\n",
        "from apache_beam.ml.rag.ingestion.mysql_common import ConflictResolution\n",
        "\n",
        "# Define how to handle conflicts - update all fields when ID matches\n",
        "conflict_resolution = ConflictResolution(\n",
        "    action=\"UPDATE\",         # Update existing records\n",
        "    update_fields=[\"embedding\", \"content\", \"metadata\"]\n",
        ")\n",
        "\n",
        "# Create writer config with conflict resolution\n",
        "cloudsql_writer_config = CloudSQLMySQLVectorWriterConfig(\n",
        "    connection_config=LanguageConnectorConfig(\n",
        "        username=DB_USER,\n",
        "        password=DB_PASSWORD,\n",
        "        database_name=DB_NAME,\n",
        "        instance_name=CONNECTION_NAME\n",
        "    ),\n",
        "    table_name=table_name,\n",
        "    conflict_resolution=conflict_resolution,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tzo43G9NfCr5"
      },
      "outputs": [],
      "source": [
        "huggingface_embedder = HuggingfaceTextEmbeddings(\n",
        "    model_name=\"sentence-transformers/all-MiniLM-L6-v2\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "axMFW_DufKnO"
      },
      "source": [
        "### Run Day 1 Pipeline\n",
        "\n",
        "First, let's ingest our initial product data:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eA3TpkHMfLUq"
      },
      "outputs": [],
      "source": [
        "# Executing on DirectRunner (local execution)\n",
        "with beam.Pipeline() as p:\n",
        "    _ = (\n",
        "        p\n",
        "        | 'Create Day 1 Products' >> beam.Create(PRODUCTS_DATA_DAY1)\n",
        "        | 'Convert Day 1 to Chunks' >> beam.Map(lambda product: Chunk(\n",
        "                content=Content(\n",
        "                    text=f\"{product['name']}: {product['description']}\"\n",
        "                ), # The text that will be embedded\n",
        "                id=product['id'],  # Use product ID as chunk ID\n",
        "                metadata=product,  # Store all product info in metadata\n",
        "            )\n",
        "          )\n",
        "        | 'Generate Day1 Embeddings' >> MLTransform(write_artifact_location=tempfile.mkdtemp())\n",
        "          .with_transform(HuggingfaceTextEmbeddings(model_name=\"sentence-transformers/all-MiniLM-L6-v2\"))\n",
        "        | 'Write Day 1 to CloudSQL MySQL' >> VectorDatabaseWriteTransform(\n",
        "            cloudsql_writer_config\n",
        "        )\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hFjtKX9tZIrI"
      },
      "source": [
        "#### Verify Initial Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lxSyaIhbZG52"
      },
      "outputs": [],
      "source": [
        "print(\"\\nAfter Day 1 ingestion:\")\n",
        "verify_embeddings_sqlalchemy(connection_name=CONNECTION_NAME, database=DB_NAME, table_name=table_name, user=DB_USER, password=DB_PASSWORD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yvOlen9qfSQ4"
      },
      "source": [
        "### Run Day 2 Pipeline\n",
        "\n",
        "Now let's process our updated product data:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r19qQs6ifVq1"
      },
      "outputs": [],
      "source": [
        "# Executing on DirectRunner (local execution)\n",
        "with beam.Pipeline() as p:\n",
        "    _ = (\n",
        "        p\n",
        "        | 'Create Day 2 Products' >> beam.Create(PRODUCTS_DATA_DAY2)\n",
        "        | 'Convert Day 2 to Chunks' >> beam.Map(lambda product: Chunk(\n",
        "                    content=Content(\n",
        "                        text=f\"{product['name']}: {product['description']}\"\n",
        "                    ), # The text that will be embedded\n",
        "                    id=product['id'],  # Use product ID as chunk ID\n",
        "                    metadata=product,  # Store all product info in metadata\n",
        "                )\n",
        "              )\n",
        "        | 'Generate Day 2 Embeddings' >> MLTransform(write_artifact_location=tempfile.mkdtemp())\n",
        "          .with_transform(HuggingfaceTextEmbeddings(model_name=\"sentence-transformers/all-MiniLM-L6-v2\"))\n",
        "        | 'Write Day 2 to CloudSQL MySQL' >> VectorDatabaseWriteTransform(\n",
        "            cloudsql_writer_config\n",
        "        )\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QbcZOIdcZWA6"
      },
      "source": [
        "#### Verify Updated Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_VpqhPAQZD4K"
      },
      "outputs": [],
      "source": [
        "print(\"\\nAfter Day 2 ingestion:\")\n",
        "verify_embeddings_sqlalchemy(connection_name=CONNECTION_NAME, database=DB_NAME, table_name=table_name, user=DB_USER, password=DB_PASSWORD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D5hImiN0fZo5"
      },
      "source": [
        "### What Changed?\n",
        "\n",
        "Key points to notice:\n",
        "\n",
        "1. The embedding vector changed because the product description was updated\n",
        "2. The metadata JSON field contains the updated category, price, and timestamp\n",
        "3. The content field reflects the new description\n",
        "4. The original ID remained the same\n",
        "\n",
        "This pattern allows you to:\n",
        "- Update embeddings when source text changes\n",
        "- Maintain referential integrity with consistent IDs\n",
        "- Track changes through the metadata field\n",
        "- Handle conflicts gracefully using CloudSQL MySQL's conflict resolution\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ndovzTB0mLdg"
      },
      "source": [
        "## Adding Embeddings to Existing Database Records <a name=\"integration\" target=\"_blank\"></a>\n",
        "\n",
        "This section demonstrates how to:\n",
        "1. Read existing product data from a database\n",
        "2. Generate embeddings for that data\n",
        "3. Write the embeddings back to the database"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l3-wl9e1fjms"
      },
      "outputs": [],
      "source": [
        "table_name = \"existing_products\"\n",
        "table_schema = \"\"\"\n",
        "    id VARCHAR(255) PRIMARY KEY,\n",
        "    title VARCHAR(255) NOT NULL,\n",
        "    description TEXT,\n",
        "    price DECIMAL,\n",
        "    embedding VECTOR(384) USING VARBINARY\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2cjjrbjUmaUN",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title MySQL helper for inserting initial records\n",
        "import sqlalchemy\n",
        "from sqlalchemy import text\n",
        "from sqlalchemy.exc import SQLAlchemyError\n",
        "# The google.cloud.sql.connector and a driver like PyMySQL (`pip install pymysql`)\n",
        "# are required for this to connect to Cloud SQL.\n",
        "from google.cloud.sql.connector import Connector\n",
        "\n",
        "# Assume get_db_engine is defined elsewhere to connect using a MySQL dialect,\n",
        "# e.g., 'mysql+pymysql'\n",
        "# from your_utils import get_db_engine\n",
        "\n",
        "def setup_initial_data_sqlalchemy(connection_name: str,\n",
        "                                  database: str,\n",
        "                                  table_name: str,\n",
        "                                  table_schema: str,\n",
        "                                  user: str,\n",
        "                                  password: str,\n",
        "                                  **connect_kwargs):\n",
        "    \"\"\"Sets up a table and inserts sample data into a MySQL database using SQLAlchemy.\n",
        "\n",
        "    This function will drop the specified table if it exists, recreate it based on the\n",
        "    provided schema, and insert a predefined set of sample products.\n",
        "\n",
        "    Args:\n",
        "        connection_name: Cloud SQL MySQL instance connection name string.\n",
        "        database: Name of the database.\n",
        "        table_name: Name of the table to create and populate.\n",
        "        table_schema: A string containing MySQL-compatible column definitions\n",
        "                      (e.g., \"id VARCHAR(255) PRIMARY KEY, title VARCHAR(255)\").\n",
        "        user: Database username.\n",
        "        password: Database password.\n",
        "        connect_kwargs: Additional keyword arguments for the Cloud SQL connector.\n",
        "    \"\"\"\n",
        "    engine = None\n",
        "    try:\n",
        "        # Assumes get_db_engine returns a SQLAlchemy engine configured for MySQL\n",
        "        engine = get_db_engine(connection_name, user, password, database, **connect_kwargs)\n",
        "\n",
        "        with engine.connect() as connection:\n",
        "            print(\"✅ Connected to Cloud SQL MySQL successfully via SQLAlchemy!\")\n",
        "\n",
        "            # DDL operations (DROP/CREATE) in MySQL cause an implicit commit,\n",
        "            # so they are run outside an explicit transaction block.\n",
        "            print(f\"Dropping table `{table_name}` if it exists...\")\n",
        "            connection.execute(text(f\"DROP TABLE IF EXISTS `{table_name}`;\"))\n",
        "\n",
        "            print(f\"Creating table `{table_name}`...\")\n",
        "            # Note: Ensure the table_schema and sample data columns match.\n",
        "            create_sql = f\"CREATE TABLE `{table_name}` ({table_schema});\"\n",
        "            connection.execute(text(create_sql))\n",
        "            print(f\"Table `{table_name}` created.\")\n",
        "\n",
        "            # Define the sample data to be inserted.\n",
        "            sample_products_dicts = [\n",
        "                 {\n",
        "                    \"id\": \"lamp-001\", \"title\": \"Artisan Table Lamp\",\n",
        "                    \"description\": \"Hand-crafted ceramic...\", \"price\": 129.99\n",
        "                 },\n",
        "                 {\n",
        "                    \"id\": \"mirror-001\", \"title\": \"Floating Wall Mirror\",\n",
        "                    \"description\": \"Modern circular mirror...\", \"price\": 199.99\n",
        "                 },\n",
        "                 {\n",
        "                     \"id\": \"vase-001\", \"title\": \"Contemporary Ceramic Vase\",\n",
        "                     \"description\": \"Minimalist vase...\", \"price\": 79.99\n",
        "                 }\n",
        "            ]\n",
        "\n",
        "            # The INSERT statement uses named placeholders matching the dictionary keys.\n",
        "            insert_sql = text(f\"\"\"\n",
        "                INSERT INTO `{table_name}` (id, title, description, price)\n",
        "                VALUES (:id, :title, :description, :price)\n",
        "            \"\"\")\n",
        "\n",
        "            print(f\"Inserting sample data into `{table_name}`...\")\n",
        "            # SQLAlchemy executes the insert for each dictionary in the list.\n",
        "            # This runs within a new transaction block started by the `connect()` context.\n",
        "            connection.execute(insert_sql, sample_products_dicts)\n",
        "\n",
        "            # Explicitly commit the transaction that contains the INSERT statements.\n",
        "            connection.commit()\n",
        "            print(\"✓ Sample products inserted successfully.\")\n",
        "\n",
        "    except SQLAlchemyError as e:\n",
        "        print(f\"❌ An SQLAlchemy error occurred during setup: {e}\")\n",
        "    except Exception as e:\n",
        "        print(f\"❌ An unexpected error occurred during setup: {e}\")\n",
        "    finally:\n",
        "        if engine:\n",
        "            # Dispose of the engine to close all connections in the pool.\n",
        "            engine.dispose()\n",
        "            print(\"Database engine pool disposed.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HjHOJ0sRmrwu"
      },
      "outputs": [],
      "source": [
        "setup_initial_data_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name, table_schema, DB_USER, DB_PASSWORD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MCN0mI08m0Ba"
      },
      "source": [
        "### Read from Database and Generate Embeddings\n",
        "\n",
        "Now let's create a pipeline to read the existing data, generate embeddings, and write back:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q2gY_kh1m08Z"
      },
      "outputs": [],
      "source": [
        "from apache_beam.io.jdbc import ReadFromJdbc\n",
        "from apache_beam.io.jdbc import WriteToJdbc\n",
        "from apache_beam.ml.rag.ingestion.mysql_common import ColumnSpecsBuilder\n",
        "\n",
        "# Configure database writer\n",
        "cloudsql_writer_config = CloudSQLMySQLVectorWriterConfig(\n",
        "    connection_config=LanguageConnectorConfig(\n",
        "        username=DB_USER,\n",
        "        password=DB_PASSWORD,\n",
        "        database_name=DB_NAME,\n",
        "        instance_name=CONNECTION_NAME\n",
        "    ),\n",
        "    table_name=table_name,\n",
        "    column_specs=(\n",
        "        ColumnSpecsBuilder()\n",
        "          .with_id_spec()\n",
        "          .with_embedding_spec()\n",
        "          # Add a placeholder value for the title column, because it has a\n",
        "          # NOT NULL constraint. Insert with Conflict resolution statements in\n",
        "          # MySQL requires all NOT NULL fields to have a value, even if the\n",
        "          # value will not be updated (the original title is preserved).\n",
        "          .add_custom_column_spec(\n",
        "            ColumnSpec.text(\"title\", value_fn=lambda x: \"\")\n",
        "           )\n",
        "          .build()\n",
        "    ),\n",
        "    conflict_resolution=ConflictResolution(\n",
        "        action=\"UPDATE\",\n",
        "        update_fields=[\"embedding\"]  # Update the embedding field\n",
        "    )\n",
        ")\n",
        "\n",
        "# Create and run pipeline  on DirectRunner (local execution)\n",
        "with beam.Pipeline() as p:\n",
        "    # Read existing products\n",
        "    rows = (\n",
        "        p\n",
        "        | \"Read Products\" >> ReadFromJdbc(\n",
        "            table_name=table_name,\n",
        "            driver_class_name=\"com.mysql.cj.jdbc.Driver\",\n",
        "            jdbc_url=cloudsql_writer_config.connector_config.to_connection_config(\n",
        "              ).jdbc_url,\n",
        "            username=DB_USER,\n",
        "            password=DB_PASSWORD,\n",
        "            query=f\"SELECT id, title, description FROM {table_name}\",\n",
        "            classpath=cloudsql_writer_config.connector_config.additional_jdbc_args()['classpath']\n",
        "        )\n",
        "    )\n",
        "\n",
        "    # Generate and write embeddings\n",
        "    _ = (\n",
        "        rows\n",
        "        | \"Convert to Chunks\" >> beam.Map(lambda row: Chunk(\n",
        "              id=row.id,\n",
        "              content=Content(text=f\"{row.title}: {row.description}\")\n",
        "            )\n",
        "          )\n",
        "        | \"Generate Embeddings\" >> MLTransform(\n",
        "            write_artifact_location=tempfile.mkdtemp()\n",
        "        ).with_transform(HuggingfaceTextEmbeddings(model_name=\"sentence-transformers/all-MiniLM-L6-v2\"))\n",
        "        | \"Write Back to CloudSQL MySQL\" >> VectorDatabaseWriteTransform(\n",
        "            cloudsql_writer_config\n",
        "        )\n",
        "    )\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bBNl5DK3Zh58"
      },
      "source": [
        "### Verify Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "elU53NLtZlTf"
      },
      "outputs": [],
      "source": [
        "print(\"\\nAfter embedding generation:\")\n",
        "verify_embeddings_sqlalchemy(connection_name=CONNECTION_NAME, database=DB_NAME, table_name=table_name, user=DB_USER, password=DB_PASSWORD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZFgpFarCp4Wo"
      },
      "source": [
        "What Happened?\n",
        "1. We started with a table containing product data but no embeddings\n",
        "2. Read the existing records using ReadFromJdbc\n",
        "3. Converted rows to Chunks, combining title and description for embedding\n",
        "4. Generated embeddings using our model\n",
        "5. Wrote back to the same table, updating only the embedding field\n",
        "Preserved all other fields (price, etc.)\n",
        "\n",
        "This pattern is useful when:\n",
        "\n",
        "- You have an existing product database\n",
        "- You want to add embeddings without disrupting current data\n",
        "- You need to maintain existing schema and relationships\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-L8mGusPd83L"
      },
      "source": [
        "## Generate Embeddings with VertexAI Text Embeddings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dVB1qAARmOlc"
      },
      "source": [
        "This section demonstrates how to use use the Vertex AI text-embeddings API to generate text embeddings that use Googles large generative artificial intelligence (AI) models.\n",
        "\n",
        "Vertex AI models are subject to [Rate Limits and Quotas](https://cloud.google.com/vertex-ai/generative-ai/docs/quotas#view-the-quotas-by-region-and-by-model) and Dataflow automatically retries throttled requests with exponential backoff.\n",
        "\n",
        "\n",
        "For more information, see [Get text embeddings](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings) in the Vertex AI documentation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-eLuZ78Tqm4w"
      },
      "source": [
        "### Authenticate with Google Cloud\n",
        "To use the Vertex AI API, we authenticate with Google Cloud."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "84p608l4ql8p"
      },
      "outputs": [],
      "source": [
        "# Replace <PROJECT_ID> with a valid Google Cloud project ID.\n",
        "PROJECT_ID = '' # @param {type:'string'}\n",
        "\n",
        "import sys\n",
        "if 'google.colab' in sys.modules:\n",
        "  from google.colab import auth\n",
        "  auth.authenticate_user(project_id=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9PZVv8S5oTHo"
      },
      "source": [
        "### Create CloudSQL MySQL table with default schema\n",
        "\n",
        "First we create a table to store our embeddings:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cEuU4JkVkLBk"
      },
      "outputs": [],
      "source": [
        "table_name = \"vertex_product_embeddings\"\n",
        "table_schema = f\"\"\"\n",
        "  id VARCHAR(255) PRIMARY KEY,\n",
        "  embedding VECTOR(768) USING VARBINARY,\n",
        "  content text,\n",
        "  metadata JSON\n",
        "\"\"\"\n",
        "setup_db_table_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name,table_schema, DB_USER, DB_PASSWORD)\n",
        "test_db_connection_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name, DB_USER, DB_PASSWORD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QZ7tSAfQpG_Z"
      },
      "source": [
        "### Configure Embedding Handler\n",
        "\n",
        "Import the `VertexAITextEmbeddings` handler, and specify the desired `textembedding-gecko` model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ipv7R6G9pqnx"
      },
      "outputs": [],
      "source": [
        "from apache_beam.ml.rag.embeddings.vertex_ai import VertexAITextEmbeddings\n",
        "\n",
        "vertexai_embedder = VertexAITextEmbeddings(model_name=\"text-embedding-005\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D7VoYav9rQJU"
      },
      "source": [
        "### Run the Pipeline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fi5SMGpZrPEm"
      },
      "outputs": [],
      "source": [
        "import tempfile\n",
        "\n",
        "# Executing on DirectRunner (local execution)\n",
        "with beam.Pipeline() as p:\n",
        "    _ = (\n",
        "            p\n",
        "            | 'Create Products' >> beam.Create(PRODUCTS_DATA)\n",
        "            | 'Convert to Chunks' >> beam.Map(lambda product: Chunk(\n",
        "                    content=Content(\n",
        "                        text=f\"{product['name']}: {product['description']}\"\n",
        "                    ), # The text that will be embedded\n",
        "                    id=product['id'],  # Use product ID as chunk ID\n",
        "                    metadata=product,  # Store all product info in metadata\n",
        "                )\n",
        "              )\n",
        "              | 'Generate Embeddings' >> MLTransform(write_artifact_location=tempfile.mkdtemp())\n",
        "                .with_transform(\n",
        "                    vertexai_embedder\n",
        "                )\n",
        "              | 'Write to CloudSQL MySQL' >> VectorDatabaseWriteTransform(\n",
        "                  CloudSQLMySQLVectorWriterConfig(\n",
        "                    connection_config=LanguageConnectorConfig(\n",
        "                        username=DB_USER,\n",
        "                        password=DB_PASSWORD,\n",
        "                        database_name=DB_NAME,\n",
        "                        instance_name=CONNECTION_NAME\n",
        "                    ),\n",
        "                    table_name=table_name\n",
        "                  )\n",
        "              )\n",
        "          )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9hVYw0rspp7Y"
      },
      "source": [
        "### Verify Embeddings"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xSEY1IILsMvi"
      },
      "outputs": [],
      "source": [
        "print(\"\\nAfter embedding generation:\")\n",
        "verify_embeddings_sqlalchemy(connection_name=CONNECTION_NAME, database=DB_NAME, table_name=table_name, user=DB_USER, password=DB_PASSWORD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yv4Rd1ZvsB_M"
      },
      "source": [
        "## Streaming Embeddings Updates from PubSub\n",
        "\n",
        "This section demonstrates how to build a real-time embedding pipeline that continuously processes product updates and maintains fresh embeddings in CloudSQL MySQL. This approach is ideal data that changes frequently.\n",
        "\n",
        "This example runs on Dataflow because streaming with DirectRunner and writing via JDBC is not supported.\n",
        "\n",
        "### Authenticate with Google Cloud\n",
        "To use the PubSub, we authenticate with Google Cloud.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VCqJmaznt1nS"
      },
      "outputs": [],
      "source": [
        "# Replace <PROJECT_ID> with a valid Google Cloud project ID.\n",
        "PROJECT_ID = '' # @param {type:'string'}\n",
        "\n",
        "import sys\n",
        "if 'google.colab' in sys.modules:\n",
        "  from google.colab import auth\n",
        "  auth.authenticate_user(project_id=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2FsoFaugtsln"
      },
      "source": [
        "### Setting Up PubSub Resources\n",
        "\n",
        "First, let's set up the necessary PubSub topics and subscriptions:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nqMe0Brlt7Bk"
      },
      "outputs": [],
      "source": [
        "from google.cloud import pubsub_v1\n",
        "from google.api_core.exceptions import AlreadyExists\n",
        "import json\n",
        "\n",
        "# Define pubsub topic\n",
        "TOPIC = \"product-updates\" # @param {type:'string'}\n",
        "\n",
        "# Create publisher client and topic\n",
        "publisher = pubsub_v1.PublisherClient()\n",
        "topic_path = publisher.topic_path(PROJECT_ID, TOPIC)\n",
        "try:\n",
        "    topic = publisher.create_topic(request={\"name\": topic_path})\n",
        "    print(f\"Created topic: {topic.name}\")\n",
        "except AlreadyExists:\n",
        "    print(f\"Topic {topic_path} already exists.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "07ZFeGbMuFj_"
      },
      "source": [
        "### Create CloudSQL MySQL Table for Streaming Updates\n",
        "\n",
        "Next, create a table to store the embedded data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3Xc70uV_uJy5"
      },
      "outputs": [],
      "source": [
        "table_name = \"streaming_product_embeddings\"\n",
        "table_schema = \"\"\"\n",
        "  id VARCHAR(255) PRIMARY KEY,\n",
        "  embedding VECTOR(384) USING VARBINARY,\n",
        "  content text,\n",
        "  metadata JSON,\n",
        "  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8HPhUfAuorBP"
      },
      "outputs": [],
      "source": [
        "setup_db_table_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name,table_schema, DB_USER, DB_PASSWORD)\n",
        "test_db_connection_sqlalchemy(CONNECTION_NAME, DB_NAME, table_name, DB_USER, DB_PASSWORD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LSriDxtsn1wH"
      },
      "source": [
        "### Configure the Pipeline options\n",
        "To run the pipeline on DataFlow we need\n",
        "- A gcs bucket for staging DataFlow files. Replace `<BUCKET_NAME>`: the name of a valid Google Cloud Storage bucket. Don't include a gs:// prefix or trailing slashes\n",
        "- Optionally set the Google Cloud region that you want to run Dataflow in. Replace `<REGION>` with the desired location\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kR0x7vzTrUlZ"
      },
      "outputs": [],
      "source": [
        "from apache_beam.options.pipeline_options import PipelineOptions, StandardOptions, SetupOptions, GoogleCloudOptions, WorkerOptions\n",
        "\n",
        "options = PipelineOptions()\n",
        "options.view_as(StandardOptions).streaming = True\n",
        "\n",
        "# Provide required pipeline options for the Dataflow Runner.\n",
        "options.view_as(StandardOptions).runner = \"DataflowRunner\"\n",
        "\n",
        "# Set the Google Cloud region that you want to run Dataflow in.\n",
        "REGION = 'us-central1' # @param {type:'string'}\n",
        "options.view_as(GoogleCloudOptions).region = REGION\n",
        "\n",
        "NETWORK = '' # @param {type:'string'}\n",
        "if NETWORK:\n",
        "  options.view_as(WorkerOptions).network = NETWORK\n",
        "\n",
        "SUBNETWORK = '' # @param {type:'string'}\n",
        "if SUBNETWORK:\n",
        "  options.view_as(WorkerOptions).subnetwork = f\"regions/{REGION}/subnetworks/{SUBNETWORK}\"\n",
        "\n",
        "options.view_as(GoogleCloudOptions).project = PROJECT_ID\n",
        "\n",
        "BUCKET_NAME = '' # @param {type:'string'}\n",
        "dataflow_gcs_location = \"gs://%s/dataflow\" % BUCKET_NAME\n",
        "\n",
        "# The Dataflow staging location. This location is used to stage the Dataflow pipeline and the SDK binary.\n",
        "options.view_as(GoogleCloudOptions).staging_location = '%s/staging' % dataflow_gcs_location\n",
        "\n",
        "# The Dataflow temp location. This location is used to store temporary files or intermediate results before outputting to the sink.\n",
        "options.view_as(GoogleCloudOptions).temp_location = '%s/temp' % dataflow_gcs_location\n",
        "\n",
        "import random\n",
        "options.view_as(GoogleCloudOptions).job_name = f\"cloudsql-streaming-embedding-ingest{random.randint(0,1000)}\"\n",
        "\n",
        "# options.view_as(SetupOptions).save_main_session = True\n",
        "options.view_as(SetupOptions).requirements_file = \"./requirements.txt\"\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gMKuccfHoDki"
      },
      "source": [
        "### Provide additional Python dependencies to be installed on Worker VM's\n",
        "\n",
        "We are making use of the HuggingFace `sentence-transformers` package to generate embeddings. Since this package is not installed on Worker VM's by default, we create a requirements.txt file with the additional dependencies to be installed on worker VM's.\n",
        "\n",
        "See [Managing Python Pipeline Dependencies](https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/) for more details.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RTGoA0SmoEvm"
      },
      "outputs": [],
      "source": [
        "!echo \"sentence-transformers\" > ./requirements.txt\n",
        "!cat ./requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eU0Sn19nqzLM"
      },
      "source": [
        "### Configure and Run Pipeline\n",
        "\n",
        "Our pipeline contains these key components:\n",
        "\n",
        "1. **Source**: Continuously reads messages from PubSub\n",
        "2. **Windowing**: Groups messages into 10-second windows for batch processing\n",
        "3. **Transformation**: Converts JSON messages to Chunk objects for embedding\n",
        "4. **ML Processing**: Generates embeddings using HuggingFace models\n",
        "5. **Sink**: Writes results to CloudSQL MySQL with conflict resolution"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w2pmJn5fqXHx"
      },
      "outputs": [],
      "source": [
        "import apache_beam as beam\n",
        "import tempfile\n",
        "import json\n",
        "\n",
        "from apache_beam.ml.transforms.base import MLTransform\n",
        "from apache_beam.ml.rag.types import Chunk, Content\n",
        "from apache_beam.ml.rag.ingestion.base import VectorDatabaseWriteTransform\n",
        "from apache_beam.ml.rag.ingestion.cloudsql import CloudSQLMySQLVectorWriterConfig\n",
        "from apache_beam.ml.rag.ingestion.cloudsql import LanguageConnectorConfig\n",
        "\n",
        "from apache_beam.ml.rag.ingestion.mysql_common import ConflictResolution\n",
        "\n",
        "from apache_beam.ml.rag.embeddings.huggingface import HuggingfaceTextEmbeddings\n",
        "from apache_beam.transforms.window import FixedWindows\n",
        "\n",
        "def parse_message(message):\n",
        "  #Parse a message containing product data.\n",
        "  product_json = json.loads(message.decode('utf-8'))\n",
        "  return Chunk(\n",
        "      content=Content(\n",
        "          text=f\"{product_json.get('name', '')}: {product_json.get('description', '')}\"\n",
        "      ),\n",
        "      id=product_json.get('id', ''),\n",
        "      metadata=product_json\n",
        "  )\n",
        "\n",
        "pipeline = beam.Pipeline(options=options)\n",
        "# Streaming pipeline\n",
        "_ = (\n",
        "    pipeline\n",
        "    | \"Read from PubSub\" >> beam.io.ReadFromPubSub(\n",
        "        topic=f\"projects/{PROJECT_ID}/topics/{TOPIC}\"\n",
        "    )\n",
        "    | \"Window\" >> beam.WindowInto(FixedWindows(10))\n",
        "    | \"Parse Messages\" >> beam.Map(parse_message)\n",
        "    | \"Generate Embeddings\" >> MLTransform(write_artifact_location=tempfile.mkdtemp())\n",
        "        .with_transform(HuggingfaceTextEmbeddings(model_name=\"sentence-transformers/all-MiniLM-L6-v2\"))\n",
        "    | \"Write to CloudSQL MySQL\" >> VectorDatabaseWriteTransform(\n",
        "        CloudSQLMySQLVectorWriterConfig(\n",
        "            connection_config=LanguageConnectorConfig(\n",
        "                username=DB_USER,\n",
        "                password=DB_PASSWORD,\n",
        "                database_name=DB_NAME,\n",
        "                instance_name=CONNECTION_NAME\n",
        "            ),\n",
        "            table_name=table_name,\n",
        "            conflict_resolution=ConflictResolution(\n",
        "                on_conflict_fields=\"id\",\n",
        "                action=\"UPDATE\",\n",
        "                update_fields=[\"embedding\", \"content\", \"metadata\"]\n",
        "            )\n",
        "        )\n",
        "    )\n",
        ")\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r7nJdc09vs98"
      },
      "source": [
        "### Create Publisher Subprocess\n",
        "The publisher simulates real-time product updates by:\n",
        "- Publishing sample product data to the PubSub topic every 5 seconds\n",
        "- Modifying prices and descriptions to represent changes\n",
        "- Adding timestamps to track update times\n",
        "- Running for 25 minutes in the background while our pipeline processes the data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "C9Bf0Nb0vY7r"
      },
      "outputs": [],
      "source": [
        "#@title Define PubSub publisher function\n",
        "import threading\n",
        "import time\n",
        "import json\n",
        "import logging\n",
        "from google.cloud import pubsub_v1\n",
        "import datetime\n",
        "import os\n",
        "import sys\n",
        "log_file = os.path.join(os.getcwd(), \"publisher_log.txt\")\n",
        "\n",
        "print(f\"Log file will be created at: {log_file}\")\n",
        "\n",
        "def publisher_function(project_id, topic):\n",
        "    \"\"\"Function that publishes sample product updates to a PubSub topic.\n",
        "\n",
        "    This function runs in a separate thread and continuously publishes\n",
        "    messages to simulate real-time product updates.\n",
        "    \"\"\"\n",
        "    time.sleep(300)\n",
        "    thread_id = threading.current_thread().ident\n",
        "\n",
        "    process_log_file = os.path.join(os.getcwd(), f\"publisher_{thread_id}.log\")\n",
        "\n",
        "    file_handler = logging.FileHandler(process_log_file)\n",
        "    file_handler.setFormatter(logging.Formatter('%(asctime)s - ThreadID:%(thread)d - %(levelname)s - %(message)s'))\n",
        "\n",
        "    logger = logging.getLogger(f\"worker.{thread_id}\")\n",
        "    logger.setLevel(logging.INFO)\n",
        "    logger.addHandler(file_handler)\n",
        "\n",
        "    logger.info(f\"Publisher thread started with ID: {thread_id}\")\n",
        "    file_handler.flush()\n",
        "\n",
        "    publisher = pubsub_v1.PublisherClient()\n",
        "    topic_path = publisher.topic_path(project_id, topic)\n",
        "\n",
        "    logger.info(\"Starting to publish messages...\")\n",
        "    file_handler.flush()\n",
        "    for i in range(300):\n",
        "        message_index = i % len(PRODUCTS_DATA)\n",
        "        message = PRODUCTS_DATA[message_index].copy()\n",
        "\n",
        "\n",
        "        dynamic_factor = 1.05 + (0.1 * ((i % 20) / 20))\n",
        "        message[\"price\"] = round(message[\"price\"] * dynamic_factor, 2)\n",
        "        message[\"description\"] = f\"PRICE UPDATE (factor: {dynamic_factor:.3f}): \" + message[\"description\"]\n",
        "\n",
        "        message[\"published_at\"] = datetime.datetime.now().isoformat()\n",
        "\n",
        "        data = json.dumps(message).encode('utf-8')\n",
        "        publish_future = publisher.publish(topic_path, data)\n",
        "\n",
        "        try:\n",
        "            logger.info(f\"Publishing message {message}\")\n",
        "            file_handler.flush()\n",
        "            message_id = publish_future.result()\n",
        "            logger.info(f\"Published message {i+1}: {message['id']} (Message ID: {message_id})\")\n",
        "            file_handler.flush()\n",
        "        except Exception as e:\n",
        "            logger.error(f\"Error publishing message: {e}\")\n",
        "            file_handler.flush()\n",
        "\n",
        "        time.sleep(5)\n",
        "\n",
        "    logger.info(\"Finished publishing all messages.\")\n",
        "    file_handler.flush()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jnUSynmjEmVr"
      },
      "source": [
        "#### Start publishing to PuBSub in background"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZnBBTwZHw7Ex"
      },
      "outputs": [],
      "source": [
        "# Launch publisher in a separate thread\n",
        "print(\"Starting publisher thread in 5 minutes...\")\n",
        "publisher_thread = threading.Thread(\n",
        "    target=publisher_function,\n",
        "    args=(PROJECT_ID, TOPIC),\n",
        "    daemon=True\n",
        ")\n",
        "publisher_thread.start()\n",
        "print(f\"Publisher thread started with ID: {publisher_thread.ident}\")\n",
        "print(f\"Publisher thread logging to file: publisher_{publisher_thread.ident}.log\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vGToqM9GoKOV"
      },
      "source": [
        "### Run Pipeline on Dataflow\n",
        "\n",
        "We launch the pipeline to run remotely on Dataflow. Once the job is launched, you can monitor its progress in the Google Cloud Console:\n",
        "1. Go to https://console.cloud.google.com/dataflow/jobs\n",
        "2. Select your project\n",
        "3. Click on the job named \"cloudsql-streaming-embedding-ingest\"\n",
        "4. View detailed execution graphs, logs, and metrics\n",
        "\n",
        "**Note**: This streaming pipeline runs indefinitely until manually stopped. Be sure to monitor usage and terminate the job in the [dataflow job console](https://console.cloud.google.com/dataflow/jobs) when finished testing to avoid unnecessary costs.\n",
        "\n",
        "### What to Expect\n",
        "After running this pipeline, you should see:\n",
        "- Continuous updates to product embeddings in the CloudSQL MySQL table\n",
        "- Price and description changes reflected in the metadata\n",
        "- New embeddings generated for updated product descriptions\n",
        "- Timestamps showing when each record was last modified"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NTibYI9rx46o"
      },
      "outputs": [],
      "source": [
        "# Run pipeline\n",
        "pipeline.run().wait_until_finish()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vX9VxJ82CTum"
      },
      "source": [
        "### Verify data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zSb1UoCSznkW"
      },
      "outputs": [],
      "source": [
        "# Verify the results\n",
        "print(\"\\nAfter embedding generation:\")\n",
        "verify_embeddings_sqlalchemy(connection_name=CONNECTION_NAME, database=DB_NAME, table_name=table_name, user=DB_USER, password=DB_PASSWORD)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "mcZATJbaOec0"
      ],
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
